import numpy
import nibabel

def interpolate_tl(data, ii, jj, kk):
    '''
    (ii, jj, kk) are given as a location in terms of the voxel coordinate system, but may not be
    exactly on the lattice. The intensity value of of the given location is returned.

    Tri-linear interpolation is used to determine the intensity. 
    (An explanation of tri-linear interpolation can be found at  http://en.wikipedia.org/wiki/Trilinear_interpolation)
    
    '''
    i_f=numpy.int(ii)
    i_c=i_f+1
    i_delta=ii-i_f

    j_f=numpy.int(jj)
    j_c=j_f+1
    j_delta=jj-j_f
   
    k_f=numpy.int(kk)
    k_c=k_f+1
    k_delta=kk-k_f

    try:
        Vi_00=data[i_f, j_f, k_f]*(1-i_delta)+data[i_c, j_f, k_f]*i_delta
    except IndexError:
        Vi_00=0

    try:
        Vi_01=data[i_f, j_f, k_c]*(1-i_delta)+data[i_c, j_f, k_c]*i_delta
    except IndexError:
        Vi_01=0

    try:
        Vi_10=data[i_f, j_c, k_f]*(1-i_delta)+data[i_c, j_c, k_f]*i_delta
    except IndexError:
        Vi_10=0

    try:
        Vi_11=data[i_f, j_c, k_c]*(1-i_delta)+data[i_c, j_c, k_c]*i_delta
    except IndexError:
        Vi_11=0

    Vj_0=Vi_00*(1-j_delta)+Vi_10*j_delta
    Vj_1=Vi_01*(1-j_delta)+Vi_11*j_delta

    #assert Vj_0*(1-k_delta)+Vj_1*k_delta >= 0

    return Vj_0*(1-k_delta)+Vj_1*k_delta

def interpolate_nn(data, ii, jj, kk):
    '''
    (ii, jj, kk) are given as a location in terms of the voxel coordinate system, but may not be
    exactly on the lattice. The intensity value of of the given location is returned.

    Nearest neighbour interpolation is used to determine the intensity. 
    '''

    # floor or ceiling of ii ??
    i_f=numpy.int(ii)
    if ii-i_f < 0.5:
        i=i_f
    else:
        i=i_f+1

    # floor and ceiling of jj ??
    j_f=numpy.int(jj)
    if jj-j_f < 0.5:
        j=j_f
    else:
        j=j_f+1

    # floor and ceiling of kk ??
    k_f=numpy.int(kk)
    if kk-k_f < 0.5:
        k=k_f
    else:
        k=k_f+1

    try:
        return data[i,j,k]
    except IndexError: # out of bound
        return 0



def world_coordinate_registration(from_image_file, to_image_file, new_image_file):
    '''
    Two (modalities of brain) image were taken in the same world coordinate. For example, patient's head
    was fixed in the imaging device (MRI), and two images were taken within a short period of time, and thus
    it can be assume that there was no head movement between these two images. (We say in this case that 
    the two images were taken in the same world coordinate. )

    But the voxel coordinat of these two image may be not aligned. For example, 
    1). the two image may have different dimension;
    2). the two image may have different resolution; or,
    3). the voxel lattice of the two images may be not aligned (translation, rotation, etc.)  

    This function transforms the voxel lattice of the "from_image" to that of "to_image" based on
    their shared world coordinate. The result is a new image that:
    1). have the same content as the "from_image" (modulo out interpolation and resampling), and
    2). have the same voxel lattice as the "old_image"

    Input: 
       -) from_image: we want content of this image.
       -) to_image:   we want the voxel lattice of this image.
    Output:
       a new image of the same content as "from_image" and the same voxel lattice as "to_image".
    '''

    # read in the from_image
    from_image=nibabel.load(from_image_file)
    #from_hdr=from_image.get_header()
    from_data=from_image.get_data()
    from_dtype=from_image.get_data_dtype()
    #from_i_spacing=from_hdr['pixdim'][1]
    #from_j_spacing=from_hdr['pixdim'][2]
    #from_k_spacing=from_hdr['pixdim'][3]
    #from_x_offset=from_hdr['qoffset_x'] 
    #from_y_offset=from_hdr['qoffset_y'] 
    #from_z_offset=from_hdr['qoffset_z']
    from_affine=from_image.get_affine()
    from_affine_inv=numpy.linalg.inv(from_affine) # world --> voxel

    # read in the to_image
    to_image=nibabel.load(to_image_file)
    to_shape=to_image.get_shape()
    #to_hdr=to_image.get_header()
    #to_data=to_image.get_data()
    #to_i_spacing=to_hdr['pixdim'][1]
    #to_j_spacing=to_hdr['pixdim'][2]
    #to_k_spacing=to_hdr['pixdim'][3]
    #to_x_offset=to_hdr['qoffset_x'] 
    #to_y_offset=to_hdr['qoffset_y'] 
    #to_z_offset=to_hdr['qoffset_z']
    to_affine=to_image.get_affine()

    # construct header for new image 
    new_hdr=to_image.get_header().copy()

    # construct data for new_image
    new_data=numpy.zeros(to_shape, from_dtype)

    # new_voxel_coordinate ==> (to_affine) world_coordinate ==> (from_affine_inv) from_image_voxel_coordinate 
    # ==> (interpolation) image intensity
    for tk in range(to_shape[2]): # slices along z-direction
        for tj in range(to_shape[1]): # in-plane along y-direction
            for ti in range(to_shape[0]): # in-plane along x-direction
                (wx,wy,wz,blah)=numpy.dot(to_affine, [ti,tj,tk,1])
                (fi,fj,fk,blah)=numpy.dot(from_affine_inv, [wx,wy,wz,1])
                intensity=interpolate_nn(from_data, fi, fj, fk)
                new_data[ti,tj,tk]=intensity.astype(from_dtype)

    # combine new_hdr and new_data to new_image, and save
    new_img=nibabel.Nifti1Image(new_data,affine=to_affine, header=new_hdr)
    new_img.to_filename(new_image_file)
